package com.redis;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.Response;

import java.io.IOException;

/**
 * zset 实现简单限流
 */
public class SimpleRateLimiter {
	private Jedis jedis;

	public SimpleRateLimiter(Jedis jedis) {
		this.jedis = jedis;
	}

	public boolean isActionAllowed(String userId, String actionKey, int period, int maxCount) throws IOException {
		String key = String.format("hist:%s:%s", userId, actionKey);
		long nowTs = System.currentTimeMillis();
		Pipeline pipe = jedis.pipelined();
		pipe.multi();
		// 增加到 zset，key 是 userID + actionKey, value 是时间戳
		pipe.zadd(key, nowTs, "" + nowTs);
		// 从当前算起，60s 之前的数据全部删除
		pipe.zremrangeByScore(key, 0, nowTs - period * 1000);
		// 统计 60s 之内的统一用户同一行为的数量
		Response<Long> count = pipe.zcard(key);
		// 重新设置超时时间 60s + 1
		pipe.expire(key, period + 1);
		pipe.exec();
		pipe.close();
		// 如果没有超过最大数量，返回 true
		return count.get() <= maxCount;
	}

	public static void main(String[] args) throws IOException {
		Jedis jedis = new Jedis();
		SimpleRateLimiter limiter = new SimpleRateLimiter(jedis);
		for (int i = 0; i < 20; i++) {
			System.out.println(limiter.isActionAllowed("laoqian", "reply", 60, 5));
		}
	}
}
